"""
R^2 = 1 - sum(y - y_pred)^2 / sum(y - y_mean)^2
"""
import numpy as np
from sklearn.metrics import r2_score


def get_data():
    """

    :return:
    """
    y_true = [3, -0.5, 2, 7]
    y_pred = [2.5, 0.0, 2, 8]

    return np.array(y_true), np.array(y_pred)


def compute_r2(y_true, y_pred):
    """

    :param y_true:
    :param y_pred:
    :return:
    """
    molecular = np.power(y_true - y_pred, 2).sum()
    denominator = np.power(y_true - np.mean(y_true), 2).sum()

    return 1 - molecular / denominator


def run():
    """
    主程序
    :return:
    """
    y_true, y_pred = get_data()

    # 自己算
    r2 = compute_r2(y_true=y_true, y_pred=y_pred)

    r2_scores = r2_score(y_true=y_true, y_pred=y_pred)

    info = 'r2_scores: my: {}, r2: {}'.format(r2, r2_scores)
    print(info)


if __name__ == '__main__':
    run()
